import os
from jupyter_client import KernelManager, find_connection_file
from queue import Empty
import re

class CodeKernel:
    def __init__(self):
        self.km = None
        self.kc = None


    def start(self):
        """Start the IPython kernel."""
        if not self.km or not self.km.is_alive():
            # Start a new kernel manager and connect to it
            self.km = KernelManager(kernel_name='my_kernel')
            self.km.start_kernel()
            self.kc = self.km.client()
            self.kc.start_channels()

            # Wait for the kernel to be ready
            self.kc.wait_for_ready()
        print('Kernel started.',self.km.is_alive())

    def shutdown(self):
        try:
            """Shutdown the IPython kernel."""
            if self.km and self.km.is_alive():
                self.kc.stop_channels()
                self.km.shutdown_kernel(now=True)
                self.km = None
                self.kc = None
        except:
            ...
    def run(self, code):
        """Execute Python code in the IPython kernel and return all output."""
        if not self.km or not self.km.is_alive():
            raise RuntimeError("Kernel is not running. Call start() first.")

        # Execute the code
        self.kc.execute(code)

        # Collect all outputs
        outputs = []
        while True:
            try:
                msg = self.kc.get_iopub_msg(timeout=30)
                msg_type = msg['msg_type']
                content = msg['content']

                if msg_type == 'stream':
                    outputs.append(content['text'])
                elif msg_type == 'execute_result' or msg_type == 'display_data':
                    if 'data' in content and 'text/plain' in content['data']:
                        outputs.append(content['data']['text/plain'])
                elif msg_type == 'error':
                    outputs.append('Error: ')
                    outputs.append(clean_traceback(content['traceback']))
                elif msg_type == 'status' and content.get('execution_state') == 'idle':
                    # Kernel is idle, we've received all output
                    break
            except Empty:
                # No more messages, assume execution is done
                break

        return '\n'.join(outputs).strip()

    def __del__(self):
        self.shutdown()

def clean_traceback(traceback_list):
    cleaned_traceback = []
    for line in traceback_list:
        # 去掉ANSI转义字符
        clean_line = re.sub(r'\x1b\[.*?m', '', line)
        cleaned_traceback.append(clean_line)
    res = '\n'.join(cleaned_traceback)
    if len(res)>1000:
        res = res[:500] + '...' + res[-500:]
    return res

def run_code(python_str, code_kernel, iscode=False):
    if iscode:
        val = code_kernel.run(python_str)
        # print('代码执行结果:', val)
        if len(val)>2000:
            val = val[:500] + '...' + val[-1500:]
        return val
    # 使用正则表达式搜索代码块
    match = re.search(r'```python(.*?)```', python_str, re.DOTALL)

    if match:
        python_str_run = match.group(1)
        # print('代码：', python_str_run)
        python_str_run = re.sub('-> (str|list|dict|int|float)\s*\n', '\n', python_str_run)

        if '<safe>' not in python_str and '假设' in python_str_run or '...' in python_str or re.search(
                '\[.*# .* \.\.\..*]',
                python_str,
                re.DOTALL):
            return f'Traceback:拒绝运行代码，因怀疑其中含有假设数据,之前的代码为：```python\n{code_kernel.cache_code}\n```\n怀疑有假设变量的代码为：{python_str_run},请重新编码，不要使用假设数据'

        val = code_kernel.run(python_str_run)

        # print('代码执行结果:', val)
        return val
    else:
        match = re.search(r'```json(.*?)```', python_str, re.DOTALL)
        if match:
            print('json结果并未保存到代码：', match.group(1))
        return '没有使用正则表达式搜索到python代码，请使用```python ```代码输出'

# 使用示例
if __name__ == "__main__":
    for i in range(100):
        kernel = CodeKernel()
        try:
            kernel.start()
            result = kernel.run("print('Hello, world!')\n1/0")
            print("Output from kernel:", result)
        finally:
            kernel.shutdown()